# -*- coding: utf-8 -*-

import rospy
import numpy as np
from math import *
from misc import config
from std_srvs.srv import SetBool
from sensor_msgs.msg import JointState
from control_msgs.msg import JointTrajectoryControllerState
from geometry_msgs.msg import PoseStamped, PoseArray, Pose

class ElfinDriver(object):
    def __init__(self):
        self.joints_pub = rospy.Publisher('elfin_basic_api/joint_goal', JointState, queue_size=1)
        self.cart_pub = rospy.Publisher('elfin_basic_api/cart_goal', PoseStamped, queue_size=1)
        self.cart_path_pub = rospy.Publisher('elfin_basic_api/cart_path_goal', PoseArray, queue_size=1)
        self.stop_client = rospy.ServiceProxy('elfin_basic_api/stop_teleop', SetBool)
        self.home_client = rospy.ServiceProxy('elfin_basic_api/home_teleop', SetBool)
        rospy.Subscriber('elfin_arm_controller/state', JointTrajectoryControllerState, self.joint_state_callback)
        self.TCP_pos = [0,0,0,0,0,0]        # 机械臂末端在基座标系下的坐标 x,y,z,r,p,y  [mm][rad]
        self.end_point_pos = [0,0,0,0,0,0]  # 笔尖在基座标系下的坐标 x,y,z,r,p,y 

    def stop(self):
        self.stop_client.call(True)

    def home(self):
        self.home_client.call(True)

    def control_joints(self, j1, j2, j3, j4, j5, j6):
        js = JointState()
        js.name = ['elfin_joint1', 'elfin_joint2', 'elfin_joint3',
         'elfin_joint4', 'elfin_joint5', 'elfin_joint6']
        js.position = [j1, j2, j3, j4, j5, j6]
        js.header.stamp = rospy.get_rostime()
        self.joints_pub.publish(js)

    def control_position(self, x=258.0, y=-31.0, z=927.0, roll=0, pitch=0, yaw=0):
        '''单位[mm][rad]'''
        if type(x) == list or type(x) == tuple:
            pos_group = x
            x = pos_group[0]
            y = pos_group[1]
            z = pos_group[2]
            roll = pos_group[3]
            pitch = pos_group[4]
            yaw = pos_group[5]
        ps = PoseStamped()
        ps.header.stamp = rospy.get_rostime()
        ps.header.frame_id = 'elfin_base_link'
        ps.pose.position.x = x/1000.0
        ps.pose.position.y = y/1000.0
        ps.pose.position.z = z/1000.0
        qw, qx, qy, qz = self._euler2quaternion([roll, pitch, yaw])
        ps.pose.orientation.x = qx
        ps.pose.orientation.y = qy
        ps.pose.orientation.z = qz
        ps.pose.orientation.w = qw
        self.cart_pub.publish(ps)

    def control_path(self, points):
        '''单位[m][rad]'''
        pa = PoseArray()
        pa.header.stamp = rospy.get_rostime()
        pa.header.frame_id = 'elfin_base_link'
        for point in points:
            ps = Pose()
            ps.position.x = point[0]/1000.0
            ps.position.y = point[1]/1000.0
            ps.position.z = point[2]/1000.0
            qw, qx, qy, qz = self._euler2quaternion(point[3:])
            ps.orientation.x = qx
            ps.orientation.y = qy
            ps.orientation.z = qz
            ps.orientation.w = qw
            pa.poses.append(ps)
        self.cart_path_pub.publish(pa)
    
    def joint_state_callback(self, msg):
        '''根据J1-J6计算默认TCP位姿和笔尖位姿'''
        joints_pos = msg.actual.positions  # 六轴位置[rad]
        T0to6 = self._T0to6Cal(joints_pos) # 正运动学计算self.TCP_pos
        self.TCP_pos = TransMartrix2xyzrpy(T0to6)
        T6to7 = xyzrpy2TransMartrix(config.BRUSH_PARAMS)
        T0to7 = T0to6.dot(T6to7)
        self.end_point_pos = TransMartrix2xyzrpy(T0to7)

    def _euler2quaternion(self, euler):
        '''欧拉角转四元数'''
        roll, pitch, yaw = euler
        cosroll = cos(roll * 0.5)
        sinroll = sin(roll * 0.5)
        cospitch = cos(pitch * 0.5)
        sinpitch = sin(pitch * 0.5)
        cosyaw = cos(yaw * 0.5)
        sinyaw = sin(yaw * 0.5)
        qw = cosroll * cosyaw * cospitch + sinroll * sinyaw * sinpitch
        qx = sinroll * cosyaw * cospitch - cosroll * sinyaw * sinpitch
        qy = cosroll * cosyaw * sinpitch + sinroll * sinyaw * cospitch
        qz = cosroll * sinyaw * cospitch - sinroll * cosyaw * sinpitch
        return (qw, qx, qy, qz)
    
    def _T0to6Cal(self, joints_pos):
        '''T0到T6的变换矩阵计算'''
        DH_params = [
            # theta,                d,      a,      alpha
            [joints_pos[0],         220,    0,      pi/2],
            [joints_pos[1]+pi/2,    0,      380,    pi],
            [joints_pos[2]+pi/2,    0,      0,      pi/2],
            [joints_pos[3],         420,    0,      -pi/2],
            [joints_pos[4],         0,      0,      pi/2],
            [joints_pos[5],         180,    0,      0]
        ]
        if config.IF_ACTUAL == True: DH_params[5][1] = 185
        T = self.DH_cal(DH_params[0])
        for params in DH_params[1:]:
            _T = self.DH_cal(params)
            T = T.dot(_T)
        return T

    def DH_cal(self, DH_params):
        '''标准DH参数法变换矩阵'''
        theta, d, a, alpha = DH_params
        martrix = np.array([
            [cos(theta),    -sin(theta)*cos(alpha), sin(theta)*sin(alpha),   a*cos(theta)],
            [sin(theta),    cos(theta)*cos(alpha),  -cos(theta)*sin(alpha), a*sin(theta)],
            [0,             sin(alpha),             cos(alpha),             d],
            [0,             0,                      0,                      1]
        ])
        return martrix
    
def TransMartrix2xyzrpy(martrix):
    '''由六轴坐标计算变换矩阵,顺序内旋zyx,外旋xyz,先旋转再平移(基座标系),正算反算应采用相对应的计算公式'''
    nx = martrix[0][0]
    ny = martrix[1][0]
    nz = martrix[2][0]
    ox = martrix[0][1]
    oy = martrix[1][1]
    ax = martrix[0][2]
    ay = martrix[1][2]
    yaw = atan2(ny, nx)
    pitch = atan2(-nz, nx*cos(yaw) + ny*sin(yaw))
    roll = atan2(-ay*cos(yaw) + ax*sin(yaw), oy*cos(yaw) - ox*sin(yaw))
    return martrix[0][3], martrix[1][3], martrix[2][3], roll, pitch, yaw

def xyzrpy2TransMartrix(xyzrpy):
    '''由六轴坐标计算变换矩阵,顺序内旋zyx,外旋xyz,先旋转再平移(基座标系),正算反算应采用相对应的计算公式'''
    x, _y, z, r, p, y = xyzrpy
    martrix = np.array([
        [cos(y)*cos(p), cos(y)*sin(p)*sin(r)-sin(y)*cos(r), cos(y)*sin(p)*cos(r)+sin(y)*sin(r), x],
        [sin(y)*cos(p), sin(y)*sin(p)*sin(r)+cos(y)*cos(r), sin(y)*sin(p)*cos(r)-cos(y)*sin(r), _y],
        [-sin(p),       cos(p)*sin(r),                      cos(p)*cos(r),                      z],
        [0,             0,                                  0,                                  1]
    ])
    return martrix